'''
Created on Aug 28, 2014

@author: Siyuan Qi
'''

import os
import kde
import numpy as np
import xml.etree.ElementTree as ET

# import matplotlib.pyplot as plt
from pymc.distributions import normal_like
from pymc.distributions import mv_normal_cov_like
import matplotlib.pyplot as plt

def loadData(imgPath, imgName):
	def getFloat(line):
		line = line.rstrip(',\r\n')
		if line != "":
			return [float(c) for c in line.split(',')]
		else:
			return None

	name = os.path.splitext(imgName)[0]
# 	orientationMap = "../../output/Images/" + name + "_orientation.bmp"
# 	segmentationMap = "../../output/Images/" + name + "_resegmentation.bmp"
	orientationMap = imgPath + name + "_orientation.bmp"
	segmentationMap = imgPath + name + "_resegmentation.bmp"

	cameraHeight = 0
	K = 0
	RT = 0
	bgPropsCount = 0
	cubPropsCount = 0
	recPropsCount = 0
	
	bgDist = BackgroundDistribution()
	cbDist = CuboidDistribution()
	
	with open(imgPath + "../" + name + "_proposals.txt") as fp:
		# Read the camera parameters
		for i, line in enumerate(fp): # Notice that i stands for the (i+1)th line
			line = line.rstrip(',\r\n')
			if i == 1:
				cameraHeight = float(line)
				bgDist.setHeight(-cameraHeight)
					
			elif i == 5:
				K = np.reshape(np.matrix(getFloat(line)), (3, 3))
			elif i == 7:
				R = np.reshape(np.matrix(getFloat(line)), (3, 3))
				R[1:3, 0:3] = - R[1:3, 0:3]
				RT = np.eye(4)
				RT[0:3, 0:3] = R
			elif i==10:
				bgPropsCount = int(line)
				break
			
		# Read the background proposals:
		for i in range(bgPropsCount):
			weight = 0
			threeD = []
			floor = []; front = []; right = []; left = []; ceiling = []
			for j, line in enumerate(fp):
				if j == 1:
					weight = float(line)
				elif j == 3:
					threeD = getFloat(line)
				elif j == 5:
					floor = getFloat(line)
				elif j == 7:
					front = getFloat(line)
				elif j == 9:
					right = getFloat(line)
				elif j == 11:
					left = getFloat(line)
				elif j == 13:
					ceiling = getFloat(line)
				elif j == 14:
					bgDist.add(threeD, weight)
					break

		# Read the Cuboid proposals:
		for i, line in enumerate(fp):
			if i == 1:
				cubPropsCount = int(line)
			if i == 2:
				break

		for i in range(cubPropsCount):
			weight = 0
			threeD = []
			for j, line in enumerate(fp):
				if j == 1:
					weight = float(line)
				elif j == 3:
					threeD = getFloat(line)
				elif j == 4:
					cbDist.addCuboid(threeD, weight)
					break

		# Read the Rectangle proposals:
		for i, line in enumerate(fp):
			if i == 1:
				recPropsCount = int(line)
			if i == 2:
				break

		recProps = []
		for i in range(recPropsCount):
			weight = 0
			twoD = []
			normal = 0
			for j, line in enumerate(fp):
				if j == 1:
					weight = float(line)
				elif j == 3:
					twoD = getFloat(line)
				elif j == 5:
					normal = int(line)
				elif j == 6:
					recProps.append(RectangleProposal(weight, twoD, normal))
					break

	return orientationMap, segmentationMap, K, RT, cameraHeight, bgDist, cbDist, recProps


class RectangleProposal():
	def __init__(self, weight, twoD, normal):
		self.weight = weight
		self.twoD = twoD
		self.normal = normal
		
class BackgroundDistribution():
	def __init__(self):
		self.backgrounds = []
		self.proposalWeights = []
		
	def setHeight(self, height):
		self.height = height

	def add(self, bg, weight):
		def getData(cuboid):
			idx = [2, 1, 0, 3]
			returnCuboid = list(cuboid)
			returnCuboid = [returnCuboid[i] for i in idx]
			returnCuboid[2] = -returnCuboid[2]
			
			for i in range(len(returnCuboid)):
				returnCuboid[i] += np.random.normal(0, 0.00001)
				
			return returnCuboid
		
		self.backgrounds.append(getData(bg[1:5]))
		self.proposalWeights.append(weight)
		
	def normalizeWeight(self, weights):
		weights = np.array(weights)
		weights -= np.ndarray.min(weights)
		weights = weights/sum(weights)
		
		return weights
	
	def fitDist(self, preferredTop = None):
		def transformWeight(weights, preferredTop):
			if preferredTop is None:
				preferredTop = 10
				
			mu = weights[0]
			tau = 1/(abs(mu - weights[preferredTop-1])/10) # TODO: Denominator to be tuned
			weights = np.array([np.exp(normal_like(w, mu, tau)) for w in weights])			
			return weights
		
		bandwidthMethod = 0.05
		values = np.array(self.backgrounds).T  # TODO: bandwidth to be tuned
		weights = self.normalizeWeight(transformWeight(np.array(self.proposalWeights), preferredTop))
		self.kernel = kde.gaussian_kde(values, weights=weights, bw_method=bandwidthMethod)
		
	
		
	def getDataProposal(self, rank = 0):
		def getCuboid(cuboid):
			returnCuboid = list(cuboid)
			returnCuboid.insert(1, self.height)
			returnCuboid.insert(1, 0)
			return returnCuboid
		
		return getCuboid(self.backgrounds[rank])
			
	def propose(self):
		def getCuboid(cuboid):
			returnCuboid = list(cuboid)
			returnCuboid.insert(1, self.height)
			returnCuboid.insert(1, 0)
			return returnCuboid
		
		sample = self.kernel.resample(1)
		pdf = self.kernel.evaluate(sample)
		return getCuboid(sample.T.tolist()[0])
	
	def evaluate(self, sample):
		dataPoint = np.array([[sample[0]], [sample[3]], [sample[4]], [sample[5]]])
		pdf = self.kernel.evaluate(dataPoint)
		return pdf.tolist()[0]
	
	def proposeForRender(self):
		def expand(cuboid):
			newCuboid = [None] * 14
			newCuboid[10:] = [2, 255, 0, 255]
			newCuboid[8] = cuboid[2]
			newCuboid[9] = cuboid[5]

			x1 = cuboid[0]
			y1 = cuboid[1]
			x2 = cuboid[3]
			y2 = cuboid[4]
			newCuboid[:8] = [x1, y1, x2, y1, x2, y2, x1, y2]
			return newCuboid

		return expand(self.propose())


class CuboidDistribution():
	def __init__(self):
		self.cuboids = []
		self.proposalWeights = []
		self.objectWeights = {}
		self.combinedWeights = {}
		self.kernels = {}
		
	def addCuboid(self, cb, weight):
		# Add cuboids from the proposals. The coordinates will be reordered
		# to be the coordinates of diagonal points. Four different possible
		# directions will be treated as four different cuboids with same
		# weight, and added to the data set. 
		def rotateCuboid(cuboid):
			idx = [[0, 1, 2, 3, 4, 5], [3, 1, 2, 0, 4, 5], [3, 4, 2, 0, 1, 5], [0, 4, 2, 3, 1, 5]]
			return [[cuboid[j] for j in idx[i]] for i in range(len(idx))]
		
		def getCuboids(cuboid):
			returnCuboid = list(cuboid)
			# Add small random number to avoid singular matrix
			returnCuboid[0] += np.random.normal(0, 0.00001)
			
			# Reorder the dimension to be (x1, y1, z1, x2, y2, z2)
			idx = [3, 1, 0, 2, 5, 4]
			returnCuboid = [returnCuboid[i] for i in idx]
			returnCuboid[1] = -returnCuboid[1]
			returnCuboid[4] = -returnCuboid[4]
			return rotateCuboid(returnCuboid)
		
		self.cuboids.extend(getCuboids(cb))
		for i in range(4):
			self.proposalWeights.append(weight) 
		
	def addObjectFilters(self, xml):
		# Read the size distributions of different objects and calculate a
		# label probability for each cuboid. The probabilities will be saved
		# in a dictionary with the key to be the name of the objects. The
		# values are initialized by the probabilities, and normalized to
		# be weights.
		def parseXML(xml):
			# Retrieve the size distribution of objects from the grammar
			xmlroot = ET.parse(xml).getroot()
			params = xmlroot.find('Distributions').findall('Parameters')
			
			filters = {}
			for param in params:
				filters[param.attrib['symbol']] = param.find('size')
				
			return filters
		
		def calculateFilterWeights(filters):
			def toFloat(strList):
				if type(strList) is list:
					return [toFloat(s) for s in strList]
				else:
					return float(strList)
				
			def scale(l, ratio):
				if type(l) is list:
					return [scale(x, ratio) for x in l]
				else:
					return l*ratio
				
			def calProb(cuboid):
				# Calculate the size prior probability
				if (cuboid[3] - cuboid[0])*(cuboid[4] - cuboid[1]) > 0:  
					size = [abs(cuboid[3]-cuboid[0]), abs(cuboid[4]-cuboid[1]), abs(cuboid[5]-cuboid[2])]
				else:
					size = [abs(cuboid[4]-cuboid[1]), abs(cuboid[3]-cuboid[0]), abs(cuboid[5]-cuboid[2])]
				pmix = 0
				for i in range(0, len(sizeProbs)):
					logp = mv_normal_cov_like(size, sizeMu[i], np.diag(sizeVar[i]))
					pmix += sizeProbs[i] * np.exp(logp)
				return pmix
				
			for symbol in filters:
				# Create a new key in the dictionary
				self.objectWeights[symbol] = []
				
				sizeMu = toFloat([mu.split() for mu in filters[symbol].get('mu').split(';')])
				sizeVar = toFloat([var.split() for var in filters[symbol].get('var').split(';')])
				sizeVar = scale(sizeVar, 2.0)
				sizeProbs = [float(p) for p in filters[symbol].get('prob').split()]
				
				for cuboid in self.cuboids:
					self.objectWeights[symbol].append(calProb(cuboid))
				
				self.objectWeights[symbol] = self.normalizeWeight(self.objectWeights[symbol])
				
		calculateFilterWeights(parseXML(xml))
		
	def normalizeWeight(self, weights):
		weights = np.array(weights)
		weights -= np.ndarray.min(weights)
		if np.ndarray.max(weights) == 0:
			weights = [1./len(weights)]*len(weights)
		else:
			weights = weights/sum(weights)

		return weights
	
	def fitDist(self, temperature, preferredTop = None):
		def transformProposalWeight(weights, preferredTop):
			if preferredTop is None:
				preferredTop = 10
				
			# TODO: Weight transform method
#  	 		mu = weights[0]
#  	 		tau = 1/np.power((abs(mu - weights[preferredTop-1])/3), 2) # TODO: Denominator to be tuned
#  	 		weights = np.array([np.exp(normal_like(w, mu, tau)) for w in weights])
			return weights

		bandwidthMethod = temperature  # TODO: bandwidth to be tuned
		values = np.array(self.cuboids).T
		proposalWeights = self.normalizeWeight(transformProposalWeight(np.array(self.proposalWeights[0:]), preferredTop))
		
		# Fit a probability distribution without reweighting by object filters. 
		self.kernels[None] = kde.gaussian_kde(values, bw_method=bandwidthMethod, weights = proposalWeights)
		self.combinedWeights[None] = np.array(proposalWeights)
		
# 		plt.plot(proposalWeights)
# 		plt.show()
		
		# Fit a probability distribution for each object.
		# Reweight each distribution using the calculated label weights.
		for symbol in self.objectWeights:
			weights = np.multiply(proposalWeights, self.objectWeights[symbol])
			weights = self.normalizeWeight(weights)
# 			print symbol
# 			plt.plot(weights)
# 			plt.show()
			self.combinedWeights[symbol] = np.array(weights)
			self.kernels[symbol] = kde.gaussian_kde(values, bw_method=bandwidthMethod, weights = weights)
	
	def propose(self, symbol=None):
		sample = self.kernels[symbol].resample(1)
		return sample.T.tolist()[0]
	
	def evaluate(self, sample, symbol=None):
		dataPoint = np.array(sample).T
		pdf = self.kernels[symbol].evaluate(dataPoint)
		return pdf.tolist()[0]
	
	def proposeForRender(self, symbol=None):
		def expand(cuboid):
			newCuboid = [None] * 14
			newCuboid[10:] = [1, 255, 0, 255]
			newCuboid[8] = cuboid[2]
			newCuboid[9] = cuboid[5]

			x1 = cuboid[0]
			y1 = cuboid[1]
			x2 = cuboid[3]
			y2 = cuboid[4]
			newCuboid[:8] = [x1, y1, x2, y1, x2, y2, x1, y2]
			return newCuboid

		return expand(self.propose(symbol))
	
	def getTopCuboids(self, symbol, topN, jump):
		def expand(cuboid):
			newCuboid = [None] * 11
			newCuboid[10] = 0
			newCuboid[8] = cuboid[2]
			newCuboid[9] = cuboid[5]

			x1 = cuboid[0]
			y1 = cuboid[1]
			x2 = cuboid[3]
			y2 = cuboid[4]
			newCuboid[:8] = [x1, y1, x2, y1, x2, y2, x1, y2]
			return newCuboid
		
		values = np.array(self.cuboids)
		
		idx = np.argsort(self.combinedWeights[symbol])[::-1]
		
		cuboids = []
		for i in range(topN):
			cuboids.append(expand(values[idx[i*jump], :]))
		return cuboids
	
	def saveFrames(self, folder, likeModel, cbCount, objCount):
		# likeModel.saveFrame(folder, self.getTopCuboids(None, 20, 4), "All")
		for symbol in self.combinedWeights:
			if symbol == None:
				likeModel.saveFrame(folder, self.getTopCuboids(None, cbCount, 4), "All")
			else:
				likeModel.saveFrame(folder, self.getTopCuboids(symbol, objCount, 2), symbol)

